///
/// Copyright (C) 2018, Cyberhaven
///
/// Permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to deal
/// in the Software without restriction, including without limitation the rights
/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
/// copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in all
/// copies or substantial portions of the Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
/// SOFTWARE.
///

#ifndef S2E_PLUGINS_PovGenerator_H
#define S2E_PLUGINS_PovGenerator_H

#include <klee/util/ExprUtil.h>
#include <s2e/S2E.h>

#include <string>

namespace s2e {
namespace plugins {
namespace pov {

const std::string VARNAME_PC = "$pc";
const std::string VARNAME_GP = "$gp";
const std::string VARNAME_ADDR = "$addr";
const std::string VARNAME_SIZE = "$size";

enum PovType { POV_GENERAL, POV_TYPE1, POV_TYPE2 };

typedef std::vector<klee::ref<klee::Expr>> ExprList;
typedef std::map<std::string /* from */, std::string /* to */> VariableRemapping;

struct PovOptions {
    PovType m_type;
    uint64_t m_faultAddress;
    uint64_t m_ipMask;
    uint64_t m_regMask;
    uint64_t m_regNum;
    size_t m_bytesBeforeSecret;
    ExprList m_extraConstraints;
    VariableRemapping m_remapping;
    PovOptions() {
        m_type = POV_GENERAL;
        m_faultAddress = 0;
        m_ipMask = 0;
        m_regMask = 0;
        m_regNum = 0;
        m_bytesBeforeSecret = 0;
    }
};

class PovGenerator : public Plugin {
protected:
    bool m_compress;

    PovGenerator(S2E *s2e) : Plugin(s2e) {
        m_compress = false;
    }

    static std::string getFileName(S2EExecutionState *state, const PovOptions &opt, const std::string &filePrefix,
                                   const std::string &fileExtWithoutDot);

    static bool writeToFile(const std::string &fileName, const void *data, unsigned size, bool compress);

    std::string writeToFile(S2EExecutionState *state, const PovOptions &opt, const std::string &filePrefix,
                            const std::string &fileExtWithoutDot, const void *data, unsigned data_size);

    bool solveConstraints(S2EExecutionState *state, const PovOptions &opt, klee::Assignment &assignment);

public:
    /** TODO: move this to some common plugin */
    static void compress(const void *in_data, size_t in_data_size, std::vector<uint8_t> &out_data);
    static std::vector<uint8_t> compress(const std::string &s);
    static std::vector<uint8_t> compress(const void *data, unsigned size);

    virtual bool generatePoV(S2EExecutionState *state, const PovOptions &opt, const std::string &filePrefix,
                             std::vector<std::string> &filePaths) = 0;
};

} // namespace pov
} // namespace plugins
} // namespace s2e

#endif
